{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1Wq4SB9A_9ic"
   },
   "source": [
    "# 🥱 LazyMergekit\n",
    "\n",
    "> 🗣️ [Large Language Model Course](https://github.com/mlabonne/llm-course)\n",
    "\n",
    "❤️ Created by [@maximelabonne](https://twitter.com/maximelabonne).\n",
    "\n",
    "This notebook allows you to easily merge multiple models using [mergekit](https://github.com/cg123/mergekit). To evaluate your merges, see [🧐 LLM AutoEval](https://colab.research.google.com/drive/1Igs3WZuXAIv9X0vwqiE90QlEPys8e8Oa?usp=sharing#scrollTo=elyxjYI_rY5W).\n",
    "\n",
    "*Special thanks to [@cg123](https://github.com/cg123) for this library and [@mrfakename](https://gist.github.com/fakerybakery) who told me about sharding (see his [Gist](https://gist.github.com/fakerybakery/d30a4d31b4f914757c1381166b9c683b)).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LGd7jlfCpNcg"
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"HermesBagel-34B-v0.1\"\n",
    "yaml_config = \"\"\"\n",
    "slices:\n",
    "  - sources:\n",
    "      - model: NousResearch/Nous-Hermes-2-Yi-34B\n",
    "        layer_range: [0, 60]\n",
    "      - model: jondurbin/bagel-dpo-34b-v0.2\n",
    "        layer_range: [0, 60]\n",
    "merge_method: slerp\n",
    "base_model: NousResearch/Nous-Hermes-2-Yi-34B\n",
    "parameters:\n",
    "  t:\n",
    "    - filter: self_attn\n",
    "      value: [0, 0.5, 0.3, 0.7, 1]\n",
    "    - filter: mlp\n",
    "      value: [1, 0.5, 0.7, 0.3, 0]\n",
    "    - value: 0.5\n",
    "dtype: bfloat16\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d5mYzDo1q96y",
    "outputId": "5136bcf3-923a-4d40-d60d-12f25eaea3bd"
   },
   "outputs": [],
   "source": [
    "# @title ## Run merge\n",
    "\n",
    "# @markdown ### Runtime type\n",
    "# @markdown Select your runtime (CPU, High RAM, GPU)\n",
    "\n",
    "runtime = \"GPU\"  # @param [\"CPU\", \"CPU + High-RAM\", \"GPU\"]\n",
    "\n",
    "# @markdown ### Mergekit arguments\n",
    "# @markdown Use the `main` branch by default, [`mixtral`](https://github.com/cg123/mergekit/blob/mixtral/moe.md) if you want to create a Mixture of Experts.\n",
    "\n",
    "branch = \"main\"  # @param [\"main\", \"mixtral\"]\n",
    "trust_remote_code = True  # @param {type:\"boolean\"}\n",
    "\n",
    "# Install mergekit\n",
    "if branch == \"main\":\n",
    "    !git clone https://github.com/cg123/mergekit.git\n",
    "    !cd mergekit && pip install -qqq -e . --progress-bar off\n",
    "elif branch == \"mixtral\":\n",
    "    !git clone -b mixtral https://github.com/cg123/mergekit.git\n",
    "    !cd mergekit && pip install -qqq -e . --progress-bar off\n",
    "    !pip install -qqq -U transformers --progress-bar off\n",
    "\n",
    "# Save config as yaml file\n",
    "with open(\"config.yaml\", \"w\", encoding=\"utf-8\") as f:\n",
    "    f.write(yaml_config)\n",
    "\n",
    "# Base CLI\n",
    "if branch == \"main\":\n",
    "    cli = \"mergekit-yaml config.yaml merge --copy-tokenizer\"\n",
    "elif branch == \"mixtral\":\n",
    "    cli = \"mergekit-moe config.yaml merge --copy-tokenizer\"\n",
    "\n",
    "# Additional arguments\n",
    "if runtime == \"CPU\":\n",
    "    cli += \" --allow-crimes --out-shard-size 1B --lazy-unpickle\"\n",
    "elif runtime == \"GPU\":\n",
    "    cli += \" --cuda --low-cpu-memory\"\n",
    "if trust_remote_code:\n",
    "    cli += \" --trust-remote-code\"\n",
    "\n",
    "print(cli)\n",
    "\n",
    "# Merge models\n",
    "!{cli}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ik0V0dF55gfU"
   },
   "outputs": [],
   "source": [
    "# @title ## Upload model to Hugging Face { display-mode: \"form\" }\n",
    "# @markdown Enter your username the name of Colab secret that stores your [Hugging Face access token](https://huggingface.co/settings/tokens).\n",
    "username = \"dfurman\"  # @param {type:\"string\"}\n",
    "token = \"HF_TOKEN\"  # @param {type:\"string\"}\n",
    "\n",
    "!pip install -qU huggingface_hub\n",
    "\n",
    "import yaml\n",
    "\n",
    "from huggingface_hub import ModelCard, ModelCardData, HfApi\n",
    "from google.colab import userdata\n",
    "from jinja2 import Template\n",
    "\n",
    "if branch == \"main\":\n",
    "    template_text = \"\"\"\n",
    "---\n",
    "license: apache-2.0\n",
    "tags:\n",
    "- merge\n",
    "- mergekit\n",
    "- lazymergekit\n",
    "{%- for model in models %}\n",
    "- {{ model }}\n",
    "{%- endfor %}\n",
    "---\n",
    "\n",
    "# {{ model_name }}\n",
    "\n",
    "{{ model_name }} is a merge of the following models using [LazyMergekit](https://colab.research.google.com/drive/1obulZ1ROXHjYLn6PPZJwRR6GzgQogxxb?usp=sharing):\n",
    "\n",
    "{%- for model in models %}\n",
    "* [{{ model }}](https://huggingface.co/{{ model }})\n",
    "{%- endfor %}\n",
    "\n",
    "## 🧩 Configuration\n",
    "\n",
    "```yaml\n",
    "{{- yaml_config -}}\n",
    "```\n",
    "\n",
    "## 💻 Usage\n",
    "\n",
    "```python\n",
    "!pip install -qU transformers accelerate\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "import transformers\n",
    "import torch\n",
    "\n",
    "model = \"{{ username }}/{{ model_name }}\"\n",
    "messages = [{\"role\": \"user\", \"content\": \"What is a large language model?\"}]\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model)\n",
    "prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "pipeline = transformers.pipeline(\n",
    "    \"text-generation\",\n",
    "    model=model,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "\n",
    "outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)\n",
    "print(outputs[0][\"generated_text\"])\n",
    "```\n",
    "\"\"\"\n",
    "\n",
    "    # Create a Jinja template object\n",
    "    jinja_template = Template(template_text.strip())\n",
    "\n",
    "    # Get list of models from config\n",
    "    data = yaml.safe_load(yaml_config)\n",
    "    if \"models\" in data:\n",
    "        models = [\n",
    "            data[\"models\"][i][\"model\"]\n",
    "            for i in range(len(data[\"models\"]))\n",
    "            if \"parameters\" in data[\"models\"][i]\n",
    "        ]\n",
    "    elif \"parameters\" in data:\n",
    "        models = [\n",
    "            data[\"slices\"][0][\"sources\"][i][\"model\"]\n",
    "            for i in range(len(data[\"slices\"][0][\"sources\"]))\n",
    "        ]\n",
    "    elif \"slices\" in data:\n",
    "        models = [\n",
    "            data[\"slices\"][i][\"sources\"][0][\"model\"] for i in range(len(data[\"slices\"]))\n",
    "        ]\n",
    "    else:\n",
    "        raise Exception(\"No models or slices found in yaml config\")\n",
    "\n",
    "    # Fill the template\n",
    "    content = jinja_template.render(\n",
    "        model_name=MODEL_NAME,\n",
    "        models=models,\n",
    "        yaml_config=yaml_config,\n",
    "        username=username,\n",
    "    )\n",
    "\n",
    "elif branch == \"mixtral\":\n",
    "    template_text = \"\"\"\n",
    "---\n",
    "license: apache-2.0\n",
    "tags:\n",
    "- moe\n",
    "- merge\n",
    "- mergekit\n",
    "- lazymergekit\n",
    "{%- for model in models %}\n",
    "- {{ model }}\n",
    "{%- endfor %}\n",
    "---\n",
    "\n",
    "# {{ model_name }}\n",
    "\n",
    "{{ model_name }} is a Mixure of Experts (MoE) made with the following models using [LazyMergekit](https://colab.research.google.com/drive/1obulZ1ROXHjYLn6PPZJwRR6GzgQogxxb?usp=sharing):\n",
    "\n",
    "{%- for model in models %}\n",
    "* [{{ model }}](https://huggingface.co/{{ model }})\n",
    "{%- endfor %}\n",
    "\n",
    "## 🧩 Configuration\n",
    "\n",
    "```yaml\n",
    "{{- yaml_config -}}\n",
    "```\n",
    "\n",
    "## 💻 Usage\n",
    "\n",
    "```python\n",
    "!pip install -qU transformers bitsandbytes accelerate\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "import transformers\n",
    "import torch\n",
    "\n",
    "model = \"{{ username }}/{{ model_name }}\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model)\n",
    "pipeline = transformers.pipeline(\n",
    "    \"text-generation\",\n",
    "    model=model,\n",
    "    model_kwargs={\"torch_dtype\": torch.float16, \"load_in_4bit\": True},\n",
    ")\n",
    "\n",
    "messages = [{\"role\": \"user\", \"content\": \"Explain what a Mixture of Experts is in less than 100 words.\"}]\n",
    "prompt = pipeline.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)\n",
    "print(outputs[0][\"generated_text\"])\n",
    "```\n",
    "\"\"\"\n",
    "\n",
    "    # Create a Jinja template object\n",
    "    jinja_template = Template(template_text.strip())\n",
    "\n",
    "    # Fill the template\n",
    "    data = yaml.safe_load(yaml_config)\n",
    "    models = [model[\"source_model\"] for model in data[\"experts\"]]\n",
    "\n",
    "    content = jinja_template.render(\n",
    "        model_name=MODEL_NAME,\n",
    "        models=models,\n",
    "        yaml_config=yaml_config,\n",
    "        username=username,\n",
    "    )\n",
    "\n",
    "# Save the model card\n",
    "card = ModelCard(content)\n",
    "card.save(\"merge/README.md\")\n",
    "\n",
    "# Defined in the secrets tab in Google Colab\n",
    "api = HfApi(token=userdata.get(token))\n",
    "\n",
    "# Upload merge folder\n",
    "api.create_repo(\n",
    "    repo_id=f\"{username}/{MODEL_NAME}\",\n",
    "    repo_type=\"model\",\n",
    "    exist_ok=True,\n",
    ")\n",
    "api.upload_folder(\n",
    "    repo_id=f\"{username}/{MODEL_NAME}\",\n",
    "    folder_path=\"merge\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z9SrqDnTfr5D"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
